package crypto

import (
	"bytes"
	"crypto/aes"
	"crypto/cipher"
	"encoding/hex"
	"errors"
	"fmt"
	"github.com/gogf/gf/v2/errors/gcode"
	"github.com/gogf/gf/v2/errors/gerror"

	"github.com/gogf/gf/v2/crypto/gaes"
)

func Test() {
	key := []byte("1234567890123456")
	res, _ := gaes.Encrypt([]byte("shit"), key)

	res2 := hex.EncodeToString(res)   // string
	res3, _ := hex.DecodeString(res2) // bytes
	fmt.Println(res3)
}

func AES_Encrypt_ECB(plain, key []byte) ([]byte, error) {
	bytCipher := make([]byte, 0)

	cipher, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	blockSize := cipher.BlockSize()
	if len(plain)%blockSize != 0 {
		return nil, errors.New("data isn't block size count!")
	}

	cipherLen := len(plain)
	bytCipher = make([]byte, cipherLen)
	for i := 0; i < cipherLen; i += blockSize {
		cipher.Encrypt(bytCipher[i:i+blockSize], plain[i:i+blockSize])
	}

	return bytCipher, nil
}
func AES_Decrypt_ECB(bytCipher, key []byte) ([]byte, error) {
	bytPlain := make([]byte, 0)

	cipher, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}

	blockSize := cipher.BlockSize()
	if len(bytCipher)%blockSize != 0 {
		return nil, errors.New("data isn't block size interger count!")
	}

	cipherLen := len(bytCipher)
	bytPlain = make([]byte, cipherLen)
	for i := 0; i < cipherLen; i += blockSize {
		cipher.Decrypt(bytPlain[i:i+blockSize], bytCipher[i:i+blockSize])
	}

	return bytPlain, nil
}
func AES_Encrypt_CBC(plantText, key []byte) ([]byte, error) {
	block, err := aes.NewCipher(key)
	if err != nil {
		return nil, err
	}
	plantText = PaddingPKCS7(plantText, block.BlockSize())

	blockModel := cipher.NewCBCEncrypter(block, key)

	ciphertext := make([]byte, len(plantText))

	blockModel.CryptBlocks(ciphertext, plantText)
	return ciphertext, nil
}
func AES_Decrypt_CBC(ciphertext, key []byte) ([]byte, error) {
	keyBytes := []byte(key)
	block, err := aes.NewCipher(keyBytes) // 选择加密算法
	if err != nil {
		return nil, err
	}

	length := len(ciphertext)
	if length%block.BlockSize() != 0 {
		return nil, errors.New("data invalid")
	}

	blockModel := cipher.NewCBCDecrypter(block, keyBytes)
	plantText := make([]byte, len(ciphertext))
	blockModel.CryptBlocks(plantText, ciphertext)
	plantText, _ = UnPaddingPKCS7(plantText, block.BlockSize())
	return plantText, nil
}

func PaddingPKCS5(src []byte, blockSize int) []byte {
	padding := blockSize - len(src)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(src, padtext...)
}
func PaddingPKCS7(src []byte, blockSize int) []byte {
	padding := blockSize - len(src)%blockSize
	padtext := bytes.Repeat([]byte{byte(padding)}, padding)
	return append(src, padtext...)
}
func UnPaddingPKCS5(src []byte, blockSize int) ([]byte, error) {
	length := len(src)
	if blockSize <= 0 {
		return nil, gerror.NewCode(gcode.CodeInvalidParameter, "invalid blocklen")
	}

	if length%blockSize != 0 || length == 0 {
		return nil, gerror.NewCode(gcode.CodeInvalidParameter, "invalid data len")
	}

	unpadding := int(src[length-1])
	if unpadding > blockSize || unpadding == 0 {
		return nil, gerror.NewCode(gcode.CodeInvalidParameter, "invalid padding")
	}

	padding := src[length-unpadding:]
	for i := 0; i < unpadding; i++ {
		if padding[i] != byte(unpadding) {
			return nil, gerror.NewCode(gcode.CodeInvalidParameter, "invalid padding")
		}
	}

	return src[:(length - unpadding)], nil
}
func UnPaddingPKCS7(src []byte, blockSize int) ([]byte, error) {
	length := len(src)
	if blockSize <= 0 {
		return nil, gerror.NewCode(gcode.CodeInvalidParameter, "invalid blocklen")
	}

	if length%blockSize != 0 || length == 0 {
		return nil, gerror.NewCode(gcode.CodeInvalidParameter, "invalid data len")
	}

	unpadding := int(src[length-1])
	if unpadding > blockSize || unpadding == 0 {
		return nil, gerror.NewCode(gcode.CodeInvalidParameter, "invalid padding")
	}

	padding := src[length-unpadding:]
	for i := 0; i < unpadding; i++ {
		if padding[i] != byte(unpadding) {
			return nil, gerror.NewCode(gcode.CodeInvalidParameter, "invalid padding")
		}
	}

	return src[:(length - unpadding)], nil
}
